Does this assert mean the MultiStepModelRunner can only be run with one step? Can you elaborate on this?

MultiStepModelRunner only takes a single step internally before returning to AsyncLLMEngine. As the multi-step is done implicitly using stateful model inputs and SequenceGroup states.

thanks for the explaination!

This is a bit confusing tho. IIRC, this was introduced by me for multi-step draft model runner? We should remove this argument and use stateful model inputs as the unify representation. Also cc @alexm-neuralmagic

yeah let's remove this argument

Let's do this in a follow up PR as it will involve spec decode as well. Will add to TODO tracker

Can we auto-correct to v2 block manager and print a warning here?

bdd3ae6

yeah let's remove this argument

Add a comment to explain why we need this variable?

c325d2c

If outputs is a part of this data structure, calling this class MutableModelInput seems confusing?

Also the current class name is probably a bit too long. Maybe something like ModelRequests? Feel free to use any other name that makes more sense here.

renamed to StatefulModelInputs. The outputs is really a cache and needed for the next step. So renamed to cached_outputs

Not a review comment, just a question: Is this op attention-backend specific?

cc @WoosukKwon

yes, eventually we will move this into attention backends API, but it may involve some refactoring to do cleanly. See some initial work done here: #7571

If we are not the last pipeline stage, why would we need to know last_sampled_token_ids? We are not running the sampler if we are not the last pipeline stage right?

Yes, however non-last PP stages need to use last_sampled_token_ids to perform in-place advance_step on GPU. And this is where we append to model_inputs so that every rank sees a consistent sampled_token_ids for the last step

@SolitaryThinker no need to block on this feedback - but if you have time - I would propose adding an example/offline_inference_multi_step.py example which instantiates an engine instance with multi-step enabled. Similar in structure to example/offline_inference.py.

An example of why this is useful - as part of the logprobs workstream, I am trying to step through the multi-step model runner with the python debugger & examine the output logprobs. I am using your multi_step/test_correctness.py in order to set up a server with multi-step enabled.

However, multi_step/test_correctness.py is an end-to-end client/server test & it is not straightforward (although technically doable) to step through the server code with the debugger because the server is in another process.

I will get around this by writing a short script which sets up an engine instance with multi-step enabled.

However, for someone else who is approaching this code for the first time, it could be helpful to have an example file (or unit test) which just sets up an engine instance with multi-step enabled and invokes inference using LLM.generate(). This could be a good way to facilitate quick debugging & also gives insight into how the server works.

Here is the offline_inference_multi_step.py script I wrote for myself to facilitate debugging, if you would like to use it.

'''
Example of setting up LLM with multi-step enabled.
In actuality, async engine would be a more sensible choice
for a real use-case. However this example is useful
for demonstration & debugging of multi-step code.
'''

from vllm import LLM, SamplingParams

# Sample prompts.
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

# Create an LLM.
llm = LLM(model="JackFram/llama-160m",
          swap_space=16,
          tensor_parallel_size=1,
          gpu_memory_utilization=0.9,
          num_scheduler_steps=8,
          use_v2_block_manager=True,
          )
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

I do wonder if there's a more generic way of doing this. If this data structure gets modified somewhere else it will not be reflected here. Maybe a loop where we check the device if the object is a tensor?

These are optionals and only set if include_gpu_probs_tensor is set in the sampler.